{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp modeling\n",
    "import os\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Body Modeling\n",
    "\n",
    "Modeling code for body model, aka BERT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "# nbdev_comment from __future__ import absolute_import, division, print_function\n",
    "\n",
    "import json\n",
    "\n",
    "import tensorflow as tf\n",
    "import transformers\n",
    "from loguru import logger\n",
    "from m3tl.params import Params\n",
    "from m3tl.utils import (get_embedding_table_from_model,\n",
    "                        get_shape_list, load_transformer_model)\n",
    "from m3tl.embedding_layer.base import DefaultMultimodalEmbedding\n",
    "\n",
    "\n",
    "class MultiModalBertModel(tf.keras.Model):\n",
    "    def __init__(self, params: Params, use_one_hot_embeddings=False):\n",
    "        super(MultiModalBertModel, self).__init__()\n",
    "        self.params = params\n",
    "        if self.params.init_weight_from_huggingface:\n",
    "            self.bert_model = load_transformer_model(\n",
    "                self.params.transformer_model_name, self.params.transformer_model_loading)\n",
    "        else:\n",
    "            self.bert_model = load_transformer_model(\n",
    "                self.params.bert_config, self.params.transformer_model_loading)\n",
    "            self.bert_model(tf.convert_to_tensor(\n",
    "                transformers.file_utils.DUMMY_INPUTS))\n",
    "        self.use_one_hot_embeddings = use_one_hot_embeddings\n",
    "\n",
    "        # multimodal input dense\n",
    "        self.embedding_layer = self.bert_model.get_input_embeddings()\n",
    "        self.multimoda_embedding = self.params.embedding_layer['model'](\n",
    "            params=self.params, embedding_layer=self.embedding_layer)\n",
    "\n",
    "    @tf.function\n",
    "    def call(self, inputs, training=False):\n",
    "        emb_inputs, embedding_tup = self.multimoda_embedding(inputs, training)\n",
    "        self.embedding_output = embedding_tup.word_embedding\n",
    "        self.model_input_mask = embedding_tup.res_input_mask\n",
    "        self.model_token_type_ids = embedding_tup.res_segment_ids\n",
    "\n",
    "        outputs = self.bert_model(\n",
    "            {'input_ids': None,\n",
    "             'inputs_embeds': self.embedding_output,\n",
    "             'attention_mask': self.model_input_mask,\n",
    "             'token_type_ids': self.model_token_type_ids,\n",
    "             'position_ids': None},\n",
    "            training=training\n",
    "        )\n",
    "        self.sequence_output = outputs.last_hidden_state\n",
    "        if 'pooler_output' in outputs:\n",
    "            self.pooled_output = outputs.pooler_output\n",
    "        else:\n",
    "            # no pooled output, use mean of token embedding\n",
    "            self.pooled_output = tf.reduce_mean(\n",
    "                outputs.last_hidden_state, axis=1)\n",
    "            outputs['pooler_output'] = self.pooled_output\n",
    "        self.all_encoder_layers = tf.stack(outputs.hidden_states, axis=1)\n",
    "        outputs = {k: v for k, v in outputs.items() if k not in (\n",
    "            'hidden_states', 'attentions')}\n",
    "        outputs['model_input_mask'] = self.model_input_mask\n",
    "        outputs['model_token_type_ids'] = self.model_token_type_ids\n",
    "        outputs['all_encoder_layers'] = self.all_encoder_layers\n",
    "        outputs['embedding_output'] = self.embedding_output\n",
    "        outputs['embedding_table'] = self.embedding_layer.weights[0]\n",
    "        return emb_inputs, outputs\n",
    "\n",
    "    def get_pooled_output(self):\n",
    "        return self.pooled_output\n",
    "\n",
    "    def get_sequence_output(self):\n",
    "        \"\"\"Gets final hidden layer of encoder.\n",
    "\n",
    "        Returns:\n",
    "          float Tensor of shape [batch_size, seq_length, hidden_size] corresponding\n",
    "          to the final hidden of the transformer encoder.\n",
    "        \"\"\"\n",
    "        return self.sequence_output\n",
    "\n",
    "    def get_all_encoder_layers(self):\n",
    "        return self.all_encoder_layers\n",
    "\n",
    "    def get_embedding_output(self):\n",
    "        \"\"\"Gets output of the embedding lookup (i.e., input to the transformer).\n",
    "\n",
    "        Returns:\n",
    "          float Tensor of shape [batch_size, seq_length, hidden_size] corresponding\n",
    "          to the output of the embedding layer, after summing the word\n",
    "          embeddings with the positional embeddings and the token type embeddings,\n",
    "          then performing layer normalization. This is the input to the transformer.\n",
    "        \"\"\"\n",
    "        return self.embedding_output\n",
    "\n",
    "    def get_embedding_table(self):\n",
    "        return get_embedding_table_from_model(self.bert_model)\n",
    "\n",
    "    def get_input_mask(self):\n",
    "        return self.model_input_mask\n",
    "\n",
    "    def get_token_type_ids(self):\n",
    "        return self.model_token_type_ids\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:bert_config not exists. will load model from huggingface checkpoint.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "Adding new problem weibo_cws, problem type: seq_tag\n",
      "Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "Adding new problem weibo_fake_cls, problem type: cls\n",
      "Adding new problem weibo_masklm, problem type: masklm\n",
      "Adding new problem weibo_pretrain, problem type: pretrain\n",
      "Adding new problem weibo_fake_regression, problem type: regression\n",
      "Adding new problem weibo_fake_vector_fit, problem type: vector_fit\n",
      "Adding new problem weibo_premask_mlm, problem type: premask_mlm\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:bert_config not exists. will load model from huggingface checkpoint.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:sampling weights: \n",
      "INFO:tensorflow:weibo_fake_cls_weibo_fake_ner: 0.3389830508474576\n",
      "INFO:tensorflow:weibo_fake_multi_cls: 0.3389830508474576\n",
      "INFO:tensorflow:weibo_masklm: 0.3220338983050847\n",
      "INFO:tensorflow:sampling weights: \n",
      "INFO:tensorflow:weibo_fake_cls_weibo_fake_ner: 0.3389830508474576\n",
      "INFO:tensorflow:weibo_fake_multi_cls: 0.3389830508474576\n",
      "INFO:tensorflow:weibo_masklm: 0.3220338983050847\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "from m3tl.test_base import TestBase\n",
    "import m3tl\n",
    "import shutil\n",
    "import numpy as np\n",
    "test_base = TestBase()\n",
    "test_base.params.assign_problem(\n",
    "    'weibo_fake_ner&weibo_fake_cls|weibo_fake_multi_cls|weibo_masklm')\n",
    "params = test_base.params\n",
    "train_dataset = m3tl.train_eval_input_fn(\n",
    "    params=params, mode=m3tl.TRAIN)\n",
    "eval_dataset = m3tl.train_eval_input_fn(\n",
    "    params=params, mode=m3tl.EVAL\n",
    ")\n",
    "\n",
    "train_dataset = train_dataset.repeat()\n",
    "\n",
    "one_batch_data = next(train_dataset.as_numpy_iterator())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`MultiModalBertModel` is transformers model with multi-modal input support. One can use it as a normal keras model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "404 Client Error: Not Found for url: https://huggingface.co/voidful/albert_chinese_tiny/resolve/main/tf_model.h5\n",
      "Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFAlbertModel: ['predictions.bias', 'predictions.dense.bias', 'predictions.LayerNorm.bias', 'predictions.LayerNorm.weight', 'predictions.decoder.weight', 'predictions.decoder.bias', 'predictions.dense.weight']\n",
      "- This IS expected if you are initializing TFAlbertModel from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing TFAlbertModel from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of TFAlbertModel were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFAlbertModel for predictions without further training.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Modal Type id mapping: \n",
      " {\n",
      "    \"class\": 0,\n",
      "    \"image\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "WARNING:tensorflow:AutoGraph could not transform <bound method Socket.send of <zmq.sugar.socket.Socket object at 0x7f6d1e289980>> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: module, class, method, function, traceback, frame, or code object was expected, got cython_function_or_method\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n",
      "WARNING: AutoGraph could not transform <bound method Socket.send of <zmq.sugar.socket.Socket object at 0x7f6d1e289980>> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: module, class, method, function, traceback, frame, or code object was expected, got cython_function_or_method\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "model = MultiModalBertModel(params=params)\n",
    "_ = model(one_batch_data)\n",
    "assert model.get_pooled_output().shape[-1] == 312\n",
    "assert len(model.get_sequence_output().shape) == 3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126a4ed0>, False), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c12685d90>, True), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126e3310>, True), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126e3790>, False), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126a4ed0>, False), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c12685d90>, True), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126e3310>, True), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126e3790>, False), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126e3310>, True), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126e3790>, False), {}).\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Found untraced functions such as embeddings_layer_call_fn, embeddings_layer_call_and_return_conditional_losses, encoder_layer_call_fn, encoder_layer_call_and_return_conditional_losses, pooler_layer_call_fn while saving (showing 5 of 115). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126a4ed0>, False), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c12685d90>, True), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126e3310>, True), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126e3790>, False), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126a4ed0>, False), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c12685d90>, True), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126e3310>, True), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126e3790>, False), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126e3310>, True), {}).\n",
      "INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 26), dtype=tf.int64, name='input_ids'), None, None, None, <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f6c126e3790>, False), {}).\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Found untraced functions such as embeddings_layer_call_fn, embeddings_layer_call_and_return_conditional_losses, encoder_layer_call_fn, encoder_layer_call_and_return_conditional_losses, pooler_layer_call_fn while saving (showing 5 of 115). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "('Not JSON Serializable:', <tf.Tensor: shape=(), dtype=int32, numpy=128>)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-5-dfeb6144af78>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[0;31m# tf.saved_model.save(serving_module, os.path.join(params.ckpt_dir, 'serving'), signatures=signatures)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[0;31m# model.save(os.path.join(params.ckpt_dir, 'serving'), signatures=model.call.get_concrete_function(spec_dict), save_traces=True)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 23\u001b[0;31m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mckpt_dir\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'serving'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     24\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m \u001b[0;31m# model.save(params.ckpt_dir, signatures=model.call.get_concrete_function(one_batch_data))\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py\u001b[0m in \u001b[0;36msave\u001b[0;34m(self, filepath, overwrite, include_optimizer, save_format, signatures, options, save_traces)\u001b[0m\n\u001b[1;32m   2000\u001b[0m     \u001b[0;31m# pylint: enable=line-too-long\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2001\u001b[0m     save.save_model(self, filepath, overwrite, include_optimizer, save_format,\n\u001b[0;32m-> 2002\u001b[0;31m                     signatures, options, save_traces)\n\u001b[0m\u001b[1;32m   2003\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2004\u001b[0m   def save_weights(self,\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py\u001b[0m in \u001b[0;36msave_model\u001b[0;34m(model, filepath, overwrite, include_optimizer, save_format, signatures, options, save_traces)\u001b[0m\n\u001b[1;32m    155\u001b[0m   \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    156\u001b[0m     saved_model_save.save(model, filepath, overwrite, include_optimizer,\n\u001b[0;32m--> 157\u001b[0;31m                           signatures, options, save_traces)\n\u001b[0m\u001b[1;32m    158\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    159\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save.py\u001b[0m in \u001b[0;36msave\u001b[0;34m(model, filepath, overwrite, include_optimizer, signatures, options, save_traces)\u001b[0m\n\u001b[1;32m     87\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mdistribution_strategy_context\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_default_replica_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     88\u001b[0m       \u001b[0;32mwith\u001b[0m \u001b[0mutils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkeras_option_scope\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msave_traces\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 89\u001b[0;31m         \u001b[0msave_lib\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfilepath\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msignatures\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptions\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     90\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     91\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0minclude_optimizer\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py\u001b[0m in \u001b[0;36msave\u001b[0;34m(obj, export_dir, signatures, options)\u001b[0m\n\u001b[1;32m   1031\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1032\u001b[0m   _, exported_graph, object_saver, asset_info = _build_meta_graph(\n\u001b[0;32m-> 1033\u001b[0;31m       obj, signatures, options, meta_graph_def)\n\u001b[0m\u001b[1;32m   1034\u001b[0m   \u001b[0msaved_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msaved_model_schema_version\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconstants\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mSAVED_MODEL_SCHEMA_VERSION\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1035\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py\u001b[0m in \u001b[0;36m_build_meta_graph\u001b[0;34m(obj, signatures, options, meta_graph_def)\u001b[0m\n\u001b[1;32m   1196\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1197\u001b[0m   \u001b[0;32mwith\u001b[0m \u001b[0msave_context\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moptions\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1198\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0m_build_meta_graph_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msignatures\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptions\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmeta_graph_def\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py\u001b[0m in \u001b[0;36m_build_meta_graph_impl\u001b[0;34m(obj, signatures, options, meta_graph_def)\u001b[0m\n\u001b[1;32m   1161\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1162\u001b[0m   object_graph_proto = _serialize_object_graph(saveable_view,\n\u001b[0;32m-> 1163\u001b[0;31m                                                asset_info.asset_index)\n\u001b[0m\u001b[1;32m   1164\u001b[0m   \u001b[0mmeta_graph_def\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mobject_graph_def\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mCopyFrom\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobject_graph_proto\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1165\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py\u001b[0m in \u001b[0;36m_serialize_object_graph\u001b[0;34m(saveable_view, asset_file_def_index)\u001b[0m\n\u001b[1;32m    753\u001b[0m   \u001b[0;32mfor\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobj_proto\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msaveable_view\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnodes\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mproto\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnodes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    754\u001b[0m     _write_object_proto(obj, obj_proto, asset_file_def_index,\n\u001b[0;32m--> 755\u001b[0;31m                         saveable_view.function_name_map)\n\u001b[0m\u001b[1;32m    756\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mproto\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    757\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py\u001b[0m in \u001b[0;36m_write_object_proto\u001b[0;34m(obj, proto, asset_file_def_index, function_name_map)\u001b[0m\n\u001b[1;32m    798\u001b[0m           version=versions_pb2.VersionDef(\n\u001b[1;32m    799\u001b[0m               producer=1, min_consumer=1, bad_consumers=[]),\n\u001b[0;32m--> 800\u001b[0;31m           metadata=obj._tracking_metadata)\n\u001b[0m\u001b[1;32m    801\u001b[0m       \u001b[0;31m# pylint:enable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    802\u001b[0m     \u001b[0mproto\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muser_object\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mCopyFrom\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mregistered_type_proto\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py\u001b[0m in \u001b[0;36m_tracking_metadata\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   3077\u001b[0m   \u001b[0;34m@\u001b[0m\u001b[0mproperty\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3078\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m_tracking_metadata\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3079\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_trackable_saved_model_saver\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtracking_metadata\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3080\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3081\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m_list_extra_dependencies_for_serialization\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mserialization_cache\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/base_serialization.py\u001b[0m in \u001b[0;36mtracking_metadata\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     53\u001b[0m     \u001b[0;31m# TODO(kathywu): check that serialized JSON can be loaded (e.g., if an\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     54\u001b[0m     \u001b[0;31m# object is in the python property)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 55\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mjson_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mEncoder\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mencode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpython_properties\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     56\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     57\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mlist_extra_dependencies_for_serialization\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mserialization_cache\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/json_utils.py\u001b[0m in \u001b[0;36mencode\u001b[0;34m(self, obj)\u001b[0m\n\u001b[1;32m     51\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     52\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mencode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 53\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mEncoder\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mencode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_encode_tuple\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     54\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     55\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/json/encoder.py\u001b[0m in \u001b[0;36mencode\u001b[0;34m(self, o)\u001b[0m\n\u001b[1;32m    197\u001b[0m         \u001b[0;31m# exceptions aren't as detailed.  The list call should be roughly\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    198\u001b[0m         \u001b[0;31m# equivalent to the PySequence_Fast that ''.join() would do.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 199\u001b[0;31m         \u001b[0mchunks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0miterencode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mo\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_one_shot\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    200\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mchunks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mlist\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    201\u001b[0m             \u001b[0mchunks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mchunks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/json/encoder.py\u001b[0m in \u001b[0;36miterencode\u001b[0;34m(self, o, _one_shot)\u001b[0m\n\u001b[1;32m    255\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkey_separator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem_separator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msort_keys\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    256\u001b[0m                 self.skipkeys, _one_shot)\n\u001b[0;32m--> 257\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0m_iterencode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mo\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    258\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    259\u001b[0m def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/json_utils.py\u001b[0m in \u001b[0;36mdefault\u001b[0;34m(self, obj)\u001b[0m\n\u001b[1;32m     48\u001b[0m       \u001b[0mitems\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mas_list\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrank\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     49\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m'class_name'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'TensorShape'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'items'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mitems\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 50\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mserialization\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_json_type\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     51\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     52\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mencode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/data/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/serialization.py\u001b[0m in \u001b[0;36mget_json_type\u001b[0;34m(obj)\u001b[0m\n\u001b[1;32m     77\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__wrapped__\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     78\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 79\u001b[0;31m   \u001b[0;32mraise\u001b[0m \u001b[0mTypeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Not JSON Serializable:'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m: ('Not JSON Serializable:', <tf.Tensor: shape=(), dtype=int32, numpy=128>)"
     ]
    }
   ],
   "source": [
    "from m3tl.run_bert_multitask import create_tensorspec_from_shape_type\n",
    "from m3tl.utils import infer_shape_and_type_from_dict\n",
    "spec_dict = create_tensorspec_from_shape_type(\n",
    "    infer_shape_and_type_from_dict(one_batch_data))\n",
    "\n",
    "\n",
    "class ServingModule(tf.Module):\n",
    "    def __init__(self):\n",
    "        super(ServingModule, self).__init__()\n",
    "        self.model = model\n",
    "\n",
    "    # @tf.function(input_signature=[v for v in spec_dict.values()])\n",
    "    def serve(self, x):\n",
    "        return self.model.call(x)\n",
    "\n",
    "\n",
    "# serving_module = ServingModule()\n",
    "# _ = serving_module.serve(one_batch_data)\n",
    "# signatures = dict(\n",
    "#     serving_default=serving_module.serve.get_concrete_function(one_batch_data)\n",
    "# )\n",
    "signatures = dict(\n",
    "    serving_default=model.call.get_concrete_function(one_batch_data)\n",
    ")\n",
    "# tf.saved_model.save(serving_module, os.path.join(params.ckpt_dir, 'serving'), signatures=signatures)\n",
    "# model.save(os.path.join(params.ckpt_dir, 'serving'), signatures=model.call.get_concrete_function(spec_dict), save_traces=True)\n",
    "model.save(os.path.join(params.ckpt_dir, 'serving'))\n",
    "\n",
    "# model.save(params.ckpt_dir, signatures=model.call.get_concrete_function(one_batch_data))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.3 64-bit ('base': conda)",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
